package com.example.demo.common;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.*;
import org.springframework.cloud.loadbalancer.core.*;
import org.springframework.http.HttpHeaders;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Random;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Created with IntelliJ IDEA 2023.
 * Description:
 * User: 故我
 * Date: 2024-02-14
 * Time: 16:13
 */
public class GrayLoadBalancer implements ReactorServiceInstanceLoadBalancer {
    private static final Log log = LogFactory.getLog(GrayLoadBalancer.class);
    private final String serviceId;
    private AtomicInteger position; //列表节点下标
    private ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider;

    public GrayLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider, String serviceId) {
        this.serviceId = serviceId;
        this.position = new AtomicInteger(new Random().nextInt(1000));
        this.serviceInstanceListSupplierProvider = serviceInstanceListSupplierProvider;
    }

    public Mono<Response<ServiceInstance>> choose(Request request) {
        ServiceInstanceListSupplier supplier = (ServiceInstanceListSupplier) this.serviceInstanceListSupplierProvider.getIfAvailable(NoopServiceInstanceListSupplier::new);
        return supplier.get(request).next().map((serviceInstances) -> {
            return this.processInstanceResponse(supplier, serviceInstances, request);
        });
    }

    private Response<ServiceInstance> processInstanceResponse(ServiceInstanceListSupplier supplier, List<ServiceInstance> serviceInstances, Request request) {
        Response<ServiceInstance> serviceInstanceResponse = this.getInstanceResponse(serviceInstances, request);
        if (supplier instanceof SelectedInstanceCallback && serviceInstanceResponse.hasServer()) {
            ((SelectedInstanceCallback) supplier).selectedServiceInstance((ServiceInstance) serviceInstanceResponse.getServer());
        }

        return serviceInstanceResponse;
    }

    private Response<ServiceInstance> getInstanceResponse(List<ServiceInstance> instances, Request request) {
        if (instances.isEmpty()) {
            if (log.isWarnEnabled()) {
                log.warn("No servers available for service: " + this.serviceId);
            }

            return new EmptyResponse();
        } else {  // 实现自定义灰度发布的 loadbalancer
            // 1、得到 request 对象，通过方法传递的方式获取参数

            // 2、从 request 对象中Header中 获取 灰度标签
            RequestDataContext requestDataContext = (RequestDataContext) request.getContext();  // 先获得请求中的内容
            HttpHeaders httpHeaders = requestDataContext.getClientRequest().getHeaders(); //获取所有的请求报头
            // 获取 值为 “gray-tag” 的头部信息值
            List<String> headerList = httpHeaders.get(GlobalVariable.GRAY_TAG);

            // 3、将得到的实例对象进行分组 【 生产服务列表 | 灰度服务列表】
            if (headerList != null && headerList.size() > 0 && headerList.get(0).equals("true")) {  // 灰度请求
                // 获取灰度列表
                List<ServiceInstance> grayList = instances.stream().filter(i -> i.getMetadata().get(GlobalVariable.GRAY_TAG) != null &&
                        i.getMetadata().get(GlobalVariable.GRAY_TAG).equals("true")).toList();
                // 判断获取到灰度列表是否为空
                if (grayList.size() > 0) {
                    instances = grayList;
                }
            } else {
                // 此时就没有灰度列表，只剩正式列表  取反就行
                instances = instances.stream().filter(i -> i.getMetadata().get(GlobalVariable.GRAY_TAG) == null ||
                        !i.getMetadata().get(GlobalVariable.GRAY_TAG).equals("true")).toList();
            }

            // 4、使用 负载均衡算法 获取到上一步列表中的某个节点（这里采用轮询）
            int pos = this.position.incrementAndGet() & Integer.MAX_VALUE;
            ServiceInstance instance = instances.get(pos % instances.size());
            return new DefaultResponse(instance);
        }
    }
}

